package com.blog.cmrpersonalblog.interceptor;

import cn.dev33.satoken.stp.StpUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

/**
 * WebSocket握手拦截器 - 用于鉴权
 */
@Slf4j
@Component
public class WebSocketAuthInterceptor implements HandshakeInterceptor {

    /**
     * 握手前拦截
     * 从请求中获取token并验证，验证通过后将userId存入WebSocket会话属性
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
            
            try {
                // 方式1: 从请求参数中获取token（推荐）
                String token = servletRequest.getServletRequest().getParameter("token");
                
                // 方式2: 从请求头中获取token（备用）
                if (token == null || token.isEmpty()) {
                    String authHeader = servletRequest.getServletRequest().getHeader("Authorization");
                    if (authHeader != null && authHeader.startsWith("Bearer ")) {
                        token = authHeader.substring(7);
                    }
                }
                
                // 验证token
                if (token == null || token.isEmpty()) {
                    log.warn("WebSocket连接失败: 缺少token");
                    return false;
                }
                
                // 使用SaToken验证token
                Object loginId = StpUtil.getLoginIdByToken(token);
                if (loginId == null) {
                    log.warn("WebSocket连接失败: token无效");
                    return false;
                }
                
                // 将userId存入WebSocket会话属性，供后续使用
                Long userId = Long.parseLong(loginId.toString());
                attributes.put("userId", userId);
                attributes.put("token", token);
                
                log.info("WebSocket握手成功: userId={}", userId);
                return true;
                
            } catch (Exception e) {
                log.error("WebSocket鉴权失败", e);
                return false;
            }
        }
        
        log.warn("WebSocket连接失败: 请求类型不支持");
        return false;
    }

    /**
     * 握手后处理
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
        if (exception != null) {
            log.error("WebSocket握手后发生异常", exception);
        }
    }
}

